import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Circle, FancyArrow

fig, ax = plt.subplots(figsize=(12, 6), dpi=300)

# 配置CVPR配色
colors = {
    "feature": "#0072B2",
    "operation": "#E69F00",
    "fusion": "#009E73",
    "gate": "#D55E00"
}

# 绘制主要处理模块
modules = [
    (0.1, 0.5, 0.15, 0.4, "Input Q", colors["feature"]),
    (0.1, 0.3, 0.15, 0.4, "Input K", colors["feature"]),
    (0.3, 0.4, 0.2, 0.6, "Channel\nAlignment", colors["operation"]),
    (0.6, 0.5, 0.2, 0.6, "Gate\nGeneration", colors["gate"]),
    (0.9, 0.5, 0.2, 0.6, "Depthwise\nFusion", colors["fusion"])
]

for x, y, w, h, label, c in modules:
    ax.add_patch(Rectangle((x-w/2, y-h/2), w, h, 
                         facecolor=c, edgecolor="k", lw=0.8))
    ax.text(x, y, label, ha="center", va="center", 
          color="white", fontsize=9, weight="bold")

# 绘制数据流箭头
arrows = [
    (0.175, 0.5, 0.25, 0.5),  # Q输入
    (0.175, 0.3, 0.25, 0.3),  # K输入
    (0.45, 0.4, 0.5, 0.5),    # 对齐后特征
    (0.7, 0.5, 0.8, 0.5),     # 门控信号
    (0.85, 0.5, 1.0, 0.5)     # 最终输出
]

for x1, y1, x2, y2 in arrows:
    ax.add_patch(FancyArrow(x1, y1, x2-x1-0.05, y2-y1, 
                          width=0.015, color=colors["feature"]))

# 绘制操作符号
ax.add_patch(Circle((0.5, 0.5), 0.03, facecolor=colors["operation"]))
ax.text(0.5, 0.5, "×", ha="center", va="center", color="white", fontsize=12)

# 美化图形
ax.set_xlim(0, 1.2)
ax.set_ylim(0, 1)
ax.axis("off")

plt.tight_layout()
plt.savefig("LightBiGatedFusion_flow.pdf", bbox_inches="tight")
plt.close()
